{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d001ba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import sys  \n",
    "import json\n",
    "import torch \n",
    "import argparse\n",
    "import numpy as np\n",
    "from PIL import Image  \n",
    "from tqdm import tqdm\n",
    "from utils import model_gen\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer  \n",
    "\n",
    "ckpt_path = 'internlm/internlm-xcomposer2-vl-7b'\n",
    "tokenizer = AutoTokenizer.from_pretrained(ckpt_path, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=\"cuda\", trust_remote_code=True).eval().cuda().to(torch.bfloat16)\n",
    "model.tokenizer = tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a127c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "lang = \"en\" # en | zh\n",
    "split = \"dev\" # dev | test\n",
    "## define split/language here ##\n",
    "\n",
    "if lang == \"en\":\n",
    "    with open(f\"data/qbench/llvisionqa_{split}.json\") as f:\n",
    "        llvqa_data = json.load(f)\n",
    "elif lang == \"zh\":\n",
    "    zh_split = \"验证集\" if split == \"dev\" else \"测试集\"\n",
    "    with open(f\"data/qbench/质衡-问答-{zh_split}.json\") as f:\n",
    "        llvqa_data = json.load(f)\n",
    "else:\n",
    "    raise NotImplementedError(\"Q-Bench does not support languages other than English (en) and Chinese (zh) yet. Contact us (https://github.com/VQAssessment/Q-Bench/) to convert  Q-Bench into more languages.\")\n",
    "\n",
    "correct = np.zeros((3,4))\n",
    "all_ = np.zeros((3,4))\n",
    "pattern = re.compile(r'[A-D]')\n",
    "answers = {}\n",
    "for llddata in tqdm((llvqa_data)):\n",
    "    t, c = llddata[\"type\"], llddata[\"concern\"]\n",
    "    if lang == \"en\":\n",
    "        message = llddata[\"question\"] \n",
    "    elif lang == \"zh\":\n",
    "        message = llddata[\"question\"] \n",
    "        \n",
    "    options_prompt = ''\n",
    "    for choice, ans in zip([\"A.\", \"B.\", \"C.\", \"D.\"], llddata[\"candidates\"]):\n",
    "        options_prompt += f\"{choice} {ans} \"\n",
    "        if \"correct_ans\" in llddata and ans == llddata[\"correct_ans\"]:\n",
    "            correct_choice = choice[0]\n",
    "            \n",
    "    text = '[UNUSED_TOKEN_146]user\\nQuestion: {}\\nContext: N/A\\nOptions: {}[UNUSED_TOKEN_145]\\n[UNUSED_TOKEN_146]assistant\\nThe answer is'.format(\n",
    "                llddata[\"question\"], options_prompt)\n",
    "    \n",
    "    img_path = f\"data/qbench/llv_dev/{split}/\" + llddata[\"img_path\"]\n",
    "    # 1st dialogue turn\n",
    "    with torch.cuda.amp.autocast(): \n",
    "        response = model_gen(model, text, img_path)\n",
    "        res = pattern.findall(response)\n",
    "        if len(res) == 0:\n",
    "            print('Error:', output_text); res = 'E'\n",
    "        else:\n",
    "            res = res[0]\n",
    "        \n",
    "     \n",
    "    llddata[\"response\"] = res\n",
    "    answers[llddata[\"img_path\"]] = res[0]\n",
    "    #print(\"[Response]: {}, [Correct Ans]: {}\".format(response, correct_choice))\n",
    "    all_[t][c] += 1\n",
    "    if res[0] not in ['A', 'B', 'C', 'D']:\n",
    "        print(\"[Response]: {}, [Correct Ans]: {}\".format(res, correct_choice))\n",
    "    if split == 'dev' and res[0] == correct_choice:\n",
    "        correct[t][c] += 1\n",
    "        \n",
    "print (correct.sum(1)/all_.sum(1))\n",
    "print (correct.sum(0)/all_.sum(0))\n",
    "print (correct.sum()/all_.sum())\n",
    "torch.save(answers, 'Output/QBench_dev_en_InternLM_XComposer_VL.json.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a66448f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
